import pyautogui as pg
import cv2,torch,win32api,win32con,win32gui,time
import numpy as np
import line,random,control,threading

time.sleep(3)
EXIT_KEY = "K"
visable = True
show_conf = 0.5
#开启控制
t1 = threading.Thread(target=control.start)
t1.start()
print("控制多线程已启动。")
#随机生成编号
filename_code = random.randint(1,99999)
data_file = open("data-{}.txt".format(filename_code),"w")

print("正在载入模型...")
model = torch.hub.load("yolov5","custom","csgo_dust2.pt",source='local')
if model != None:
    print("模型载入成功！")
else:
    input("模型载入失败！")
    exit()

for i in range(3):
    print("{}秒后开始".format(3-i))
    time.sleep(1)

while True:
    time1 = time.time()
    #hot key
    if win32api.GetAsyncKeyState(ord(EXIT_KEY)):
        data_file.close()
        print("数据已经保存。")
        exit()
    #找到窗口
    hwd = win32gui.FindWindow(0,"Counter-Strike: Global Offensive - Direct3D 9")
    #判断hwd是否有效
    if hwd == 0:
        input("没有找到游戏窗口，请重新尝试！")
        exit()
    window_bbox = win32gui.GetWindowRect(hwd)
    #截图
    img = pg.screenshot(region=window_bbox)
    #推理
    results = model(img).pandas().xyxy[0]
    #处理结果
    xmins = results["xmin"]
    ymins = results["ymin"]
    xmaxs = results["xmax"]
    ymaxs = results["ymax"]
    confs = results["confidence"]
    names = results["name"]
    classes = results["class"]
    if visable:
        #图片处理
        img = np.array(img)
        img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
        #绘制方框
        for xmin,ymin,xmax,ymax,conf,name in zip(xmins, ymins, xmaxs, ymaxs, confs, names):
            #满足设定的置信度
            if conf >= show_conf:
                mouse_pos = win32api.GetCursorPos()
                img = cv2.rectangle(img,(int(xmin),int(ymin)),(int(xmax),int(ymax)),(255,0,0),2)
                img = cv2.putText(img,name+"|"+str(conf),(int(xmin),int(ymin)),cv2.FONT_HERSHEY_SIMPLEX,1,(255,0,0))
        #显示
        img = cv2.resize(img,(640,480))
        cv2.resizeWindow("showing",640,480)
        cv2.imshow("showing",img)
        if cv2.waitKey(25) & 0xFF == ord(EXIT_KEY):
            cv2.destroyAllWindows()
            break
    
    #记录
    #收集最近的10个对象信息
    count = 0
    objects = []
    for xmin,ymin,xmax,ymax,conf,name,obj_class in zip(xmins, ymins, xmaxs, ymaxs, confs, names,classes):
        #满足特定conf
        if conf < show_conf:
            continue
        if count >= 10:
            break
        #计算中心点
        cen_x = xmin + (xmax - xmin)/2
        cen_y = ymin + (ymax - ymin)/2
        objects.append([cen_x,cen_y,obj_class,conf])
        count += 1
    #确保有10个对象
    if len(objects) < 10:
        #添加
        for i in range(10-len(objects)):
            objects.append([0,0,0,0])   #添加o数据
    #记录移动操作
    kb = 0 #0w 1s 2a 3d 4未移动
    if win32api.GetAsyncKeyState(ord("W")):
        kb = 0
    elif win32api.GetAsyncKeyState(ord("S")):
        kb = 1
    elif win32api.GetAsyncKeyState(ord("A")):
        kb = 2
    elif win32api.GetAsyncKeyState(ord("D")):
        kb = 3
    else:
        kb = 4
    #记录视角操作
    view = 0 #0不动，1左边，2右边
    #左
    if win32api.GetAsyncKeyState(37):
        view = 1
    #右
    elif win32api.GetAsyncKeyState(39):
        view = 2
    else:
        view = 0
    #记录射击操作
    shot = 0 #0射击 1不射击
    if win32api.GetAsyncKeyState(96):
        shot = 0
    else:
        shot = 1
    #保存到文件
    data = ""
    #拼接对象数据
    for i in objects:
        data += str(i[0])+","+str(i[1])+","+str(i[2])+","+str(i[3])+","
    #拼接操作
    data += str(kb)+","+str(view)+","+str(shot)+ "\n"
    data_file.write(data)
    print("用时：{}s".format(time.time() - time1))